#include <BWTA.h>
#include <ScoutManager.h>
#include <UnitGroupManager.h>
#include <math.h>
#include "../BasicAIModule/Source/BasicAIModule.h"
#include "../BasicAIModule/Source/ThreatMap.h"
#include "../BasicAIModule/Source/Potential.h"

std::pair<std::list<BWTA::BaseLocation*>, double> getBestPathHelper(std::set<BWTA::BaseLocation* > baseLocations)
{
  std::pair<std::list<BWTA::BaseLocation*>, double> shortest_path;
  shortest_path.second=0;
  if (baseLocations.empty()) return shortest_path;
  if (baseLocations.size()==1)
  {
    shortest_path.first.push_back(*baseLocations.begin());
    return shortest_path;
  }
  for(std::set<BWTA::BaseLocation*>::iterator i=baseLocations.begin();i!=baseLocations.end();i++)
  {
    BWTA::BaseLocation* node=*i;
    std::set<BWTA::BaseLocation* > baseLocations2=baseLocations;
    baseLocations2.erase(*i);
    std::pair<std::list<BWTA::BaseLocation*>, double> path_result=getBestPathHelper(baseLocations2);
    double dist=path_result.second+node->getGroundDistance(path_result.first.front());
    if (dist<shortest_path.second || shortest_path.first.empty())
    {
      path_result.first.push_front(node);
      shortest_path=std::make_pair(path_result.first,dist);
    }
  }
  return shortest_path;
}

std::pair<std::list<BWTA::BaseLocation*>, double> getBestPath(std::set<BWTA::BaseLocation* > baseLocations)
{
  std::pair<std::list<BWTA::BaseLocation*>, double> shortest_path;
  shortest_path.second=0;
  if (baseLocations.empty()) return shortest_path;
  BWTA::BaseLocation* start=BWTA::getStartLocation(BWAPI::Broodwar->self());
  baseLocations.erase(start);
  if (baseLocations.size()==1)
  {
    shortest_path.first.push_back(*baseLocations.begin());
    return shortest_path;
  }
  for(std::set<BWTA::BaseLocation*>::iterator i=baseLocations.begin();i!=baseLocations.end();i++)
  {
    BWTA::BaseLocation* node=*i;
    std::set<BWTA::BaseLocation* > baseLocations2=baseLocations;
    baseLocations2.erase(*i);
    std::pair<std::list<BWTA::BaseLocation*>, double> path_result=getBestPathHelper(baseLocations2);
    double dist=start->getGroundDistance(node)+node->getGroundDistance(path_result.first.front())+path_result.second;
    if (dist<shortest_path.second || shortest_path.first.empty())
    {
      path_result.first.push_front(node);
      shortest_path=std::make_pair(path_result.first,dist);
    }
  }
  return shortest_path;
}
ScoutManager::ScoutManager(Arbitrator::Arbitrator<BWAPI::Unit*,double> *arbitrator)
{
  this->arbitrator = arbitrator;
  this->initialized = false;
  this->initialize();
}

void ScoutManager::initialize()
{
  // Stop if we are not analyzed
  if (!BasicAIModule::analyzed) return;

  myStartLocation = BWTA::getStartLocation(BWAPI::Broodwar->self());
  std::set<BWTA::BaseLocation *> startLocations = BWTA::getStartLocations();
  startLocations.erase(myStartLocation);
  std::set<BWTA::BaseLocation *>::iterator l_next;
  for(std::set<BWTA::BaseLocation *>::iterator l = startLocations.begin(); l != startLocations.end(); l=l_next)
  {
    l_next=l;
    l_next++;
    if (myStartLocation->getGroundDistance(*l) <= 0)
    {
      startLocations.erase(*l);
    }
  }

  std::list<BWTA::BaseLocation*> path=getBestPath(startLocations).first;
  for(std::list<BWTA::BaseLocation*>::iterator p=path.begin();p!=path.end();p++)
    baseLocationsToScout.push_back(*p);
  baseLocationsExplored.insert(myStartLocation);
  this->initialized = true;
}

void ScoutManager::onOffer(std::set<BWAPI::Unit*> units)
{
  // Check our job queue
  std::set<BWAPI::Unit*>::iterator u;
  std::set<BWAPI::Unit*>::iterator u_next;
  for(u = units.begin(); u != units.end(); u=u_next) {
    u_next = u;
    u_next++; // helper iterator in case we erase the element that u is pointing to and cannot ++

    // Check if this unit is in our job list
    UnitType uType = (*u)->getType();
    if (scoutJobs.find(uType) != scoutJobs.end() && scoutJobs[uType].size() > 0) {
      // Scan list
      for(std::list<ScoutPriority>::iterator job=scoutJobs[uType].begin(); job != scoutJobs[uType].end(); ++job){
        if (job->status == ScoutPriority::Triggered) {
          // accept and assign
          job->status = ScoutPriority::Assigned;
          arbitrator->accept(this, *u);

          ScoutData temp;
          temp.target = NULL;
          temp.lastTarget = NULL;
          scouts[*u] = temp;
          // TODO: Assign job->target in scout data

          // initialize PheramoneTrail for unit
          //pheramones.insert(std::make_pair(*u,new PheramoneTrail(*u,5,Potential::epsilon*1.5)));

          units.erase(u);
          break;
        }
      }
    }
  }

  //decline remaining units
  for(u = units.begin(); u != units.end(); u++)
  {
    arbitrator->decline(this, *u, 0);
  }
}

void ScoutManager::onRevoke(BWAPI::Unit *unit, double bid)
{
  onRemoveUnit(unit);
}

void ScoutManager::update()
{
  // If we haven't initialized yet, we need to try and do it now
  if (!this->initialized){
    this->initialize();
    
    // If that was not successful, return
    if (!this->initialized)
      return;
  }

  // Check our future scout jobs
  for(std::map<BWAPI::UnitType, std::list<ScoutPriority> >::iterator unitJobs=scoutJobs.begin(); unitJobs != scoutJobs.end(); ++unitJobs){
    // This unit type has a list to check through
    if (unitJobs->second.size() > 0){
      for(std::list<ScoutPriority>::iterator job=unitJobs->second.begin(); job != unitJobs->second.end(); ++job){
        // Dormant jobs
        if (job->status == ScoutPriority::Dormant && job->foodCount <= BWAPI::Broodwar->self()->supplyUsed()){
          // Flag this job
          job->status = ScoutPriority::Triggered;
        }

        // Triggered jobs
        if (job->status == ScoutPriority::Triggered) {
          // Bid on units
          std::set<BWAPI::Unit*> usefulUnits = SelectAll(unitJobs->first)(isCompleted).not(isCarryingMinerals,isCarryingGas,isGatheringGas);
          arbitrator->setBid(this, usefulUnits, job->priority);
        }
      }
    }
  }

  if (spyingScouts.size() > 0 ){
  // motivates Spying units to hang around for a while, unless a threat is present
    std::map<BWAPI::Unit*, int>::const_iterator itr;
    for (itr = spyingScouts.begin(); itr != spyingScouts.end();){
      BWAPI::Unit* unit  = (*itr).first;
      int      startTime = (*itr).second;
      bool threatPresent = false;
      bool timeOut       = false;
      // is there an immediate threat?
      std::set<BWAPI::Unit*> antiAirUnits = 
        SelectAllEnemy()(Marine, Hydralisk, Dragoon, Photon_Cannon).onlyNearestBaseLocation(scouts[unit].target);
      if (antiAirUnits.size() > 0) {
        Broodwar->printf("Scout::Threat detected!");
        threatPresent = true;
      }
      // waits around for at most 20s 
      if ( Broodwar->getFrameCount() - startTime > 24*100 ) timeOut = true;

      if (threatPresent){
        // release this scout immediately from spying role
        Broodwar->printf("Scout done spying.");
        scouts[unit].mode=ScoutData::Idle;
        scouts[unit].lastTarget = scouts[unit].target;
        spyingScouts.erase(itr++);
      }
      else 
        ++itr;
    }
  }

  if (!scouts.empty()) {
    updateScoutAssignments();
  }

  if (debugMode)
  {
    drawAssignments();
  }

  if (!scouts.empty()) {
    // Tell scouts where to move based on threat map
    ThreatMap* map = ThreatMap::create();
    InformationManager* info = InformationManager::create();
    std::list<Potential*> potentials;
    // If the enemy has the potential to attack the Overlord
    if (Broodwar->enemy()->getRace() == BWAPI::Races::Terran || // Terran have Barracks too quickly, so always assume
        // info->enemyHasBuilt(BWAPI::UnitTypes::Terran_Barracks) || 
        info->enemyHasBuilt(BWAPI::UnitTypes::Zerg_Hydralisk_Den) || 
        info->enemyHasBuilt(BWAPI::UnitTypes::Protoss_Cybernetics_Core))
    {
      // Create potentials keeping units out of sight of the enemy (for overlords only; drone scout does not see far enough to do this)
      for each (Unit* e in TheInformationManager->getAllKnownEnemy()) {
        // padding the range because it seems they can see further than the circle that is drawn
        int enemySight = (int)(info->getType(e).sightRange() * 1.05);
        Position pos = info->getLastPosition(e);
        Potential p(1.0,enemySight,&pos);
        potentials.push_back(&p);

        Broodwar->drawCircle(CoordinateType::Map,pos.x(),pos.y(),enemySight,Colors::Yellow,false);
      }
    }
    int maxSightRange = UnitTypes::Terran_Siege_Tank_Tank_Mode.sightRange(); // Make attractive force have further range than any unit's vision
    double K = -Potential::epsilon*1.1;
    for each (std::pair<Unit*,ScoutData> scout in scouts) {
      if (scout.second.mode != ScoutData::Idle) {
        Position pos = scout.second.target->getPosition();
        Potential target(K,maxSightRange,&pos); // Used to attract scout from distance
        Position* moveTo;
        //std::list<Potential*> trail = pheramones[scout.first]->getTrail();
        if (scout.first->getType().isFlyer()) {
          potentials.push_back(&target);
          // Code for pheramone trail to prevent units getting stuck in local minima.  Whenever this code runs, the Overlord just flies
          // straight to the enemy Nexus as if none of the sight potentials were there, so I have given up on this for now.

          //if (trail.size() > 0) {
          //  for each(Potential* pheramone in trail) {
          //    potentials.push_back(pheramone);
          //  }
          //}
          moveTo = map->calculateMove(scout.first,potentials);
          potentials.remove(&target);
          //if (trail.size() > 0) {
          //  for each(Potential* pheramone in trail) {
          //    potentials.remove(pheramone);
          //  }
          //}
        }
        else {
          std::list<Potential*> targetOnly; // for units with short vision range
          targetOnly.push_back(&target);
          //if (trail.size() > 0) {
          //  for each(Potential* pheramone in pheramones[scout.first]->getTrail()) {
          //    targetOnly.push_back(pheramone);
          //  }
          //}
          moveTo = map->calculateMove(scout.first,targetOnly);
        }
        //if (Broodwar->getFrameCount()%5 == 0) pheramones[scout.first]->update();
        scout.first->move(moveTo->makeValid());

        Broodwar->drawCircle(CoordinateType::Map,pos.x(),pos.y(),maxSightRange,Colors::Green,false);
      }
    }
  }
}
void ScoutManager::setInformationManager(InformationManager* infoManager)
{
  /* empty function */
}

std::string ScoutManager::getName() const
{
  return "Scout Manager";
}

std::string ScoutManager::getShortName() const
{
  return "Scout";
}

void ScoutManager::onRemoveUnit(BWAPI::Unit* unit)
{
  if (scouts.find(unit) != scouts.end())
  {
    BWTA::BaseLocation* lostTarget = scouts[unit].target;
    if (scouts[unit].mode != ScoutData::Idle && baseLocationsExplored.find(lostTarget) == baseLocationsExplored.end())
    {
      // The target needed to be scouted and was not reached
      // This selector will find all enemy buildings within range of the base
      UnitGroup enemyUnits = SelectAllEnemy().onlyNearestBaseLocation(lostTarget);
      if (enemyUnits.size() > 0){
        // Consider it scouted and occupied
        baseLocationsToScout.remove(lostTarget);
        baseLocationsExplored.insert(lostTarget);
      } else {
        // Unsure if occupied
        baseLocationsToScout.push_back(lostTarget);
      }
      // Mark as scouted regardless to avoid frequent visits
      TheInformationManager->setLastScouted(lostTarget, Broodwar->getFrameCount());

      if (debugMode)
      {
        BWAPI::Broodwar->printf("Reassigning (%d,%d)", lostTarget->getPosition().x(), lostTarget->getPosition().y());
      }
    }
    scouts.erase(unit);

    // erase spy data
    spyingScouts.erase(unit);
  }
}

void ScoutManager::setDebugMode(bool debugMode)
{
  this->debugMode = debugMode;
}

void ScoutManager::drawAssignments()
{
  //draw target vector for each scout
  for (std::map<BWAPI::Unit*,ScoutData>::iterator s = scouts.begin(); s != scouts.end(); s++)
  {
    if ((*s).second.mode != ScoutData::Idle)
    {
      BWAPI::Position scoutPos = (*s).first->getPosition();
      BWAPI::Position targetPos = (*s).second.target->getPosition();
      BWAPI::Broodwar->drawLineMap(scoutPos.x(), scoutPos.y(), targetPos.x(), targetPos.y(), BWAPI::Colors::Yellow);
      BWAPI::Broodwar->drawCircleMap(scoutPos.x(), scoutPos.y(), 6, BWAPI::Colors::Yellow);
      BWAPI::Broodwar->drawCircleMap(targetPos.x(), targetPos.y(), (*s).first->getType().sightRange(), BWAPI::Colors::Yellow);
    }
  }
}

bool ScoutManager::isScoutingBase(BWTA::BaseLocation* base) 
{
  std::map<BWAPI::Unit*, ScoutData>::iterator u;
  for(u = scouts.begin(); u != scouts.end(); u++) {
    if (u->second.mode == ScoutData::Searching && 
        u->second.target == base) {
      return true;
    }
  }

  return false;
}

int ScoutManager::numScouts() 
{
  // return non-spying scouts
  int count = 0;
  std::map<BWAPI::Unit*, ScoutData>::iterator u;
  for(u = scouts.begin(); u != scouts.end(); u++) {
    if (u->second.mode == ScoutData::Searching) {
      count++;
    }
  }

  return count;
}

void ScoutManager::addScout(BWAPI::Unit* u)
{
  ScoutData temp;
  scouts.insert(std::make_pair(u,temp));
}

bool baseCompare(BWTA::BaseLocation* a, BWTA::BaseLocation* b) { 
  return TheInformationManager->getLastScouted(a) < TheInformationManager->getLastScouted(b); 
}

void ScoutManager::updateScoutAssignments()
{
  // Remove scout positions if the enemy is not there.
  std::map<BWAPI::Unit*, ScoutData>::iterator u;
  std::set<BWAPI::Unit*> idleScouts; // group idle scouts for later
  for(u = scouts.begin(); u != scouts.end(); u++)
  {
    if ( (*u).second.mode == ScoutData::Searching
      && (*u).first->getPosition().getDistance((*u).second.target->getPosition()) < (*u).first->getType().sightRange())
    {
      bool empty = true;
      for(int x=(*u).second.target->getTilePosition().x();x<(*u).second.target->getTilePosition().x()+4;x++)
      {
        for(int y=(*u).second.target->getTilePosition().y();y<(*u).second.target->getTilePosition().y()+3;y++)
        {
          std::set<BWAPI::Unit*> unitsOnTile = BWAPI::Broodwar->getUnitsOnTile(x,y);
          for each(BWAPI::Unit* u in unitsOnTile)
          {
            if (u->getType().isResourceDepot() && u->getPlayer()->isEnemy(BWAPI::Broodwar->self()))
            {
              empty = false;
              break;
            }
          }
          if (!empty) break;
        }
        if (!empty) break;
      }
      if (empty)
      {
        TheInformationManager->setBaseEmpty((*u).second.target);
      }
      BWTA::BaseLocation* exploredBaseLocation  = (*u).second.target;
      baseLocationsToScout.remove(exploredBaseLocation);
      baseLocationsExplored.insert(exploredBaseLocation);
      TheInformationManager->setLastScouted(exploredBaseLocation, Broodwar->getFrameCount());
      
      // TONG: now makes this unit stop for a fixed duration, or until a threat is detected
      // whichever is earlier
      if (!empty) {
        if (debugMode)
        {
          BWAPI::Broodwar->printf("Successfully scouted (%d,%d). Now spying...", exploredBaseLocation->getPosition().x(), exploredBaseLocation->getPosition().y());
        }
        (*u).second.mode = ScoutData::Spying;
        this->spyingScouts[(*u).first] = Broodwar->getFrameCount();
        arbitrator->setBid(this, (*u).first, 25);
      }
      else { // go on to the next task
        (*u).second.mode = ScoutData::Idle;
        (*u).second.lastTarget = (*u).second.target;
      }
    }

    // Check if the base was taken by us, no need to scout then
    if ( (*u).second.mode == ScoutData::Searching &&
         TheBaseManager->getBase((*u).second.target) ){
      // Mark as scouted and set idle
      TheInformationManager->setLastScouted((*u).second.target, BWAPI::Broodwar->getFrameCount());
      (*u).second.mode = ScoutData::Idle;
    }

    if ((*u).second.mode == ScoutData::Idle) {
      idleScouts.insert(u->first);
    }
  }

  // Skip if we don't have scouts
  if (idleScouts.empty()) {
    return;
  }

  // Set scouts to scout.
  // Check if we have found our enemy's main yet
  std::list<BaseLocation*> priorityTargets;
  std::list<BaseLocation*> secondaryTargets;
  BWTA::BaseLocation* enemyStart = NULL;
  for each (BaseLocation* base in TheInformationManager->getEnemyBases()) {
    if (base->isStartLocation()) {
      enemyStart = base;
      break;
    }
  }
  // Prioritize start locations if we haven't found the enemy in the beginning
  // If he lost his main base but we've seen an expansion, we don't want to look
  // at he start locations anymore
  if (!enemyStart && TheInformationManager->getEnemyBases().empty()) {
    for each (BaseLocation* base in BWTA::getStartLocations()) {
      // Only push if it isn't ours
      if (TheInformationManager->getLastScouted(base) == 0 && !TheBaseManager->getBase(base) && !isScoutingBase(base)) {
        priorityTargets.push_back(base);
      }
    }
  }                   
  // if we know which one it is but have not scouted it (2 player maps)
  if (enemyStart && !TheInformationManager->getLastScouted(enemyStart) && !isScoutingBase(enemyStart)) { 
    priorityTargets.push_back(enemyStart);
  }
  
  // Push the rest of the bases, I guess?
  for each (BaseLocation* base in BWTA::getBaseLocations()) {
    // If it's not ours or the enemy's
    if (!TheBaseManager->getBase(base) && !isScoutingBase(base)) {
      secondaryTargets.push_back(base);
    }
  }
  // Sort by last scouted
  secondaryTargets.sort(baseCompare);
  // Add them onto priority, if any
  for each (BaseLocation* base in secondaryTargets) {
    priorityTargets.push_back(base);
  }
  // Pass out bases to scouts
  for each (Unit* u in idleScouts) {
    // Only proceed if this is an idle scout and we have bases to possibly scout
    if (scouts[u].mode == ScoutData::Idle && !priorityTargets.empty()) {
      BaseLocation* target = priorityTargets.front();
      // If the base exists and another scout is not heading there and this base
      // wasn't our last target (to prevent repeated scout-find-release-repeat sequences)
      if (target != NULL && !isScoutingBase(target) && target != scouts[u].lastTarget) {
        scouts[u].mode = ScoutData::Searching;
        u->rightClick(target->getPosition());
        scouts[u].target = target;
        baseLocationsToScout.remove(target);
        if (debugMode) {
          BWAPI::Broodwar->printf("Scouting (%d,%d)", target->getPosition().x(), target->getPosition().y());
        }
      }
      priorityTargets.pop_front();
    }

    // Release if still idle
    if (scouts[u].mode == ScoutData::Idle) {
      this->arbitrator->removeBid(this, u);
    }
  }
}

void ScoutManager::addJob(int priority, int foodCount, BWAPI::UnitType unit, BWAPI::TilePosition target)
{
  ScoutPriority temp;
  temp.priority = priority;
  temp.foodCount = foodCount;
  temp.target = target;

  // Default unit type is a worker
  if (unit == BWAPI::UnitTypes::None){
    unit = BWAPI::Broodwar->self()->getRace().getWorker();
  }

  this->scoutJobs[unit].push_back(temp);
}
